# 第一步：导入必要的库和配置环境
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import get_image_colorizer
import torch
import warnings
import gradio as gr
from PIL import Image

# 关闭fastai的一些警告信息
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")

# 设置设备为GPU0，如果有多个GPU，可以选择GPU1至GPU7
device.set(device=DeviceId.GPU0)

# 检查GPU是否可用
if not torch.cuda.is_available():
    print('GPU not available. Switching to CPU...')
    device.set(device=DeviceId.CPU)

# 加载DeOldify模型，不指定model_path
colorizer = get_image_colorizer(artistic=True)

# 第二步：定义图像处理函数
def colorize_image(input_img, render_factor, watermarked):
    # 上色处理
    colorized_img = colorizer.get_transformed_image(input_img, render_factor=render_factor, watermarked=watermarked)
    
    # 将 PIL Image 转换为 RGB 模式（如果尚未转换）
    if colorized_img.mode != "RGB":
        colorized_img = colorized_img.convert("RGB")
    
    return colorized_img

# 第三步：创建Gradio接口
def gradio_interface():
    # 定义Gradio界面
    interface = gr.Interface(
        fn=colorize_image,  # 调用的处理函数
        inputs=[
            gr.Image(type="pil", label="上传图像"),  # 上传黑白图像
            gr.Slider(7, 40, step=1, value=26, label="渲染系数 "),  # 渲染因子
            
        ],
        outputs=gr.Image(type="pil", label="Colorized Image"),  # 输出彩色图像
        title="图像上色器",  # 应用标题
        description="上传黑白照片以通过DeOldify 为其着色。调整渲染因子以获得更好的效果。",  # 应用描述
        theme="default"  # Gradio主题
    )
    
    # 启动应用
    interface.launch(share=True)

# 运行Gradio界面
if __name__ == "__main__":
    gradio_interface()
